/*
 * bpp_sequential.cc
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * Block Principal Pivoting, solving for local minimum of
 *   L = sum |A - W H^T|^2
 *   subject to W >= 0, H >= 0
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/mpi.hpp>

#include "helper/helper.h"
#include "algorithm/StopWatch.h"
#include "data/FileAccess.h"
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"
#include "arma_helper.h"

class BlockPrincipalPivoting
{
public:
  // input/output variables
  arma::mat A, W, H;
  // iteration information
  std::vector<Iteration> iterInfo;
protected:
  // algorithm variables
  int m, n, k;
  double l1, l2;
  int maxiter, interval;
  bool normalizing;
  int bigiter;
public:
  BlockPrincipalPivoting() { 
    PO_REQUIRE(vm, "iter");
    PO_REQUIRE(vm, "k");

    PO_VARIABLE_VALUE(k, vm, "k", int);
    l1 = (vm.count("l1")) ? vm["l1"].as<double>() : 0; // l1 regularization
    l2 = (vm.count("l2")) ? vm["l2"].as<double>() : 0; // l2 regularization
    PO_VARIABLE_VALUE(maxiter, vm, "iter", int);
    PO_VARIABLE_VALUE(interval, vm, "interval", int);
    normalizing = vm.count("normalize");
  }

  void LoadData() 
  {
    PO_REQUIRE(vm, "data");
    PO_VARIABLE(data, vm, "data", std::string);
    
    using namespace arma;
    std::stringstream ss;
    ss << data << "_" << MPI_RANK;
    A.load(ss.str(), arma_binary);
    m = A.n_rows;
    n = A.n_cols;
  }

  static void ThresholdVector(arma::vec& v)
  {
    for (unsigned int i = 0; i < v.n_elem; i++)
      if (v(i) < 1e-16) v(i) = 0;
  }

  static void AbsThresholdVector(arma::vec& v)
  {
    for (unsigned int i = 0; i < v.n_elem; i++)
      if (fabs(v(i)) < 1e-16) v(i) = 0;
  }

  static void NegativeIndex(const arma::vec& x, const arma::uvec& F, arma::uvec& H)
  {
    //MSG("x=" << arma::trans(x) << "\nF=" << arma::trans(F));
    int count = 0;
    for (unsigned int i=0; i < F.n_elem; i++)
      if (x(F(i)) < 0) count++;
    H.zeros(count);
    int j = 0;
    for (unsigned int i=0; i < F.n_elem; i++)
      if (x(F(i)) < 0) H(j++) = F(i);
  }

  static void SetValue(arma::vec& x, const arma::uvec& F, const arma::vec& x_F)
  {
    x.zeros(x.n_elem);
    for (unsigned int i=0; i < F.n_elem; i++)
      x(F(i)) = x_F(i);
  }

  static void SetMarker(arma::uvec& marker, const arma::uvec& F, bool set)
  {
    for (unsigned int i=0; i < F.n_elem; i++)
      marker(F(i)) = set;
  }

  static void GetMarker(arma::uvec& marker, arma::uvec& F)
  {
    int count = 0;
    for (unsigned int i=0; i < marker.n_elem; i++)
      if (marker(i)) count++;
    F.zeros(count);
    int j = 0;
    for (unsigned int i=0; i < marker.n_elem; i++)
      if (marker(i)) F(j++) = i;
  }

  void ExchangeIndex(int q, arma::uvec& F, const arma::uvec& H1, arma::uvec& G, const arma::uvec& H2)
  {
    arma::uvec marker(q);
    
    // F = F-H1+H2
    marker.zeros(q);
    SetMarker(marker, F, true);
    SetMarker(marker, H2, true);
    SetMarker(marker, H1, false);
    GetMarker(marker, F);

    // G= G-H2+H1
    marker.zeros(q);
    SetMarker(marker, G, true);
    SetMarker(marker, H1, true);
    SetMarker(marker, H2, false);
    GetMarker(marker, G);    
  }

  void SolveNNLS(const arma::mat& C, const arma::vec& b, arma::vec& x)
  {
    int pbar = 3;
    int p = pbar;
    int q = C.n_cols;
    unsigned int t = q + 1;
    
    using namespace arma;
    uvec F, G(q);
    for (int i = 0; i < q; i++) G(i) = i;
    x.zeros(q);
    arma::vec y = -trans(C)*b;

    //MSG(C.n_rows << " " << C.n_cols << " " << x.n_elem << " " << y.n_elem);
    
    for (int iter = 0; iter < 10*q; iter++) {
      uvec H1, H2;

      NegativeIndex(x, F, H1);
      NegativeIndex(y, G, H2);

      //MSG("F=" << F.n_elem << " H1=" << H1.n_elem << " G=" << G.n_elem << " H2=" << H2.n_elem);
      if (H1.n_elem + H2.n_elem == 0) { // success
	//MSG("success");
	return;
      }
      else if ( H1.n_elem + H2.n_elem < t) {
	t = H1.n_elem + H2.n_elem;
	p = pbar;
	ExchangeIndex(q, F, H1, G, H2);
      }
      else if (p >= 1) {
	p = p-1;
	ExchangeIndex(q, F, H1, G, H2);
      }
      else {
	int maxindex = -1;
	bool from_F = true;
	for (unsigned int i = 0; i < H1.n_elem; i++)
	  if (maxindex < (int)H1(i)) maxindex = H1(i);
	for (unsigned int i = 0; i < H2.n_elem; i++)
	  if (maxindex < (int)H2(i)) { maxindex = H2(i); from_F = false; }
	if (from_F) {
	  uvec h1(1), h2;
	  h1(0) = maxindex;
	  ExchangeIndex(q, F, h1, G, h2);
	}
	else {
	  uvec h1, h2(1);
	  h2(0) = maxindex;
	  ExchangeIndex(q, F, h1, G, h2);
	}
      }

      //arma::vec x_F = solve(C.cols(F), b);
      arma::vec x_F = solve(C.submat(F,F), b.elem(F));
      AbsThresholdVector(x_F);
      //arma::vec y_G = trans(C.cols(G)) * (C.cols(F)*x_F - b);
      arma::vec y_G = C.submat(G,F)*x_F - b.elem(G);
      AbsThresholdVector(y_G);

      SetValue(x, F, x_F);
      SetValue(y, G, y_G);
    }
    MSG("failure");
  }

  void SolveBPP()
  {
    StopWatch watch, clock;
    TimeDuration runtime;

    using namespace arma;
    srand(100);
    W = randu<mat>(m, k); 
    H = randu<mat>(n, k); 

    for (bigiter = 0; bigiter < maxiter; bigiter++) {
      clock.GetDuration();
      // update W
      arma::mat CC = trans(H)*H;        // size k x k
      arma::mat CA = trans(H)*trans(A); // size k x m
      for (int i = 0; i < m; i++) { // solve for i-th row of W (size 1 x k)
	//MSG("w" << i);
	arma::vec w_i/*, b*/;
	//b = trans(A.row(i));        // size b n x 1
	//SolveNNLS(H, b, w_i);       // H size n x k, w_i size k x 1
	SolveNNLS(CC, CA.col(i), w_i);  // 
	W.row(i) = trans(w_i);
      }

      // update H
      CC = trans(W)*W; // k x k
      CA = trans(W)*A; // k x n
      for (int i = 0; i < n; i++) { // solve for i-th row of H (size 1 x k)
	//MSG("h" << i);
	arma::vec h_i/*, b*/;
	//b = A.col(i);               // size b m x 1
	// SolveNNLS(W, b, h_i);       // W size m x k, h_i size k x 1
	SolveNNLS(CC, CA.col(i), h_i);
	H.row(i) = trans(h_i);
      }
      runtime += clock.GetDuration();

      if (bigiter % interval == 0 || bigiter == maxiter-1) {
        double res = norm(A-W*trans(H), "fro");
        if (MPI_RANK == 0) {
          MSG("it=" << bigiter << " |res|=" << res  
	      << " |W|=" << norm(W,"fro") 
	      << " |H|=" << norm(H,"fro")
              << " d=" << watch.GetDuration());
	  Iteration info(bigiter, res);
	  info.duration = watch.GetTotalDuration();
	  info.compute = runtime;
	  iterInfo.push_back(info);
	}
      }
    }
    MSG(watch.GetTotalDuration());
    MSG("compute clock " << runtime);
  }
};

void run_bpp()
{
  if (MPI_RANK == 0) MSG("Block Principal Pivoting NMF");
  BlockPrincipalPivoting a;
  a.LoadData();
  a.SolveBPP();
  
  if (MPI_RANK == 0 && vm.count("output")) {
    PO_VARIABLE(output, vm, "output", std::string);
    a.H.save(output, arma::raw_ascii);
  }

  if (MPI_RANK == 0 && vm.count("iterinfo")) {
    PO_VARIABLE(infoFile, vm, "iterinfo", std::string);
    FileAccessWrite(out, infoFile, NULL);
    const std::vector<Iteration>& info =  a.iterInfo;
    for (unsigned int i = 0; i < info.size(); i++) {
      out << info[i].iter << "\t" << (double)info[i].duration.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].compute.total_milliseconds()/1000 << "\t" 
	  << (double)info[i].communication.total_milliseconds()/1000 << "\t"
	  << info[i].residual << std::endl;
    }
  }
}
